import numpy as np


# 状态转移模型：x_k = x_(k-1) + v_k
def state_transition(x, noise_std=1.0):
    return x + np.random.normal(0, noise_std, size=x.shape)


# 观测模型：z_k = x_k + w_k
def observation_model(x, observation, noise_std=1.0):
    return np.exp(-0.5 * ((observation - x) / noise_std) ** 2) / (np.sqrt(2 * np.pi) * noise_std)


# 粒子滤波函数
def particle_filter(num_particles, initial_state, observations, transition_std=1.0, obs_std=1.0):
    particles = np.random.normal(initial_state, 1.0, size=num_particles)  # 初始化粒子
    weights = np.ones(num_particles) / num_particles  # 初始权重

    estimates = []

    for z in observations:
        # 预测
        particles = state_transition(particles, noise_std=transition_std)

        # 更新权重
        weights *= observation_model(particles, z, noise_std=obs_std)
        weights /= np.sum(weights)  # 归一化

        # 重采样
        indices = np.random.choice(np.arange(num_particles), size=num_particles, p=weights)
        particles = particles[indices]
        weights = np.ones(num_particles) / num_particles

        # 状态估计
        estimate = np.mean(particles)
        estimates.append(estimate)

    return estimates


# 示例：跟踪状态的变化
true_states = np.linspace(0, 10, 50)
observations = true_states + np.random.normal(0, 1.0, size=true_states.shape)

estimates = particle_filter(num_particles=1000, initial_state=0, observations=observations)


# 可视化结果
import matplotlib.pyplot as plt

# 可视化结果
plt.plot(true_states, label="True State")
plt.plot(observations, label="Observations", alpha=0.5)
plt.plot(estimates, label="Particle Filter Estimate")
plt.legend()
plt.title("Particle Filter Tracking with Increased Particles")
plt.xlabel("Time Step")
plt.ylabel("State Value")
plt.show()